from flask import Flask, request, jsonify, render_template_string
import requests
import os

app = Flask(__name__)

# 配置参数（建议使用环境变量）
# app.config["API_KEY"] = os.getenv("ANYTHINGLLM_API_KEY", "Bearer YOUR_API_KEY")
# app.config["API_BASE"] = os.getenv("ANYTHINGLLM_API_BASE", "http://localhost:3001/api")
# app.config["WORKSPACE_ID"] = os.getenv("WORKSPACE_ID", "")  # 可选工作区ID

API_KEY = "Bearer E84CN7G-G974QK6-JZE6QNP-3PPN0XD"
API_BASE = "http://localhost:3001/api/v1"


# WORKSPACE_ID = "test"  # 可选工作区ID


@app.route('/upload', methods=['POST'])
def upload_file():
    """文件上传接口"""
    # 验证文件存在
    if 'file' not in request.files:
        return jsonify({"error": "No file uploaded"}), 400

    file = request.files['file']
    if file.filename == '':
        return jsonify({"error": "Empty filename"}), 400

    # 调用AnythingLLM API
    try:
        response = requests.post(
            url=f"{API_BASE}/document/upload/",
            headers={"Authorization": API_KEY},
            files={"file": (file.filename, file.stream, file.mimetype)},
            data={"name": request.form.get('name', file.filename),
                  # "workspaceId": WORKSPACE_ID
                  },
            timeout=30
        )
        response.raise_for_status()
        return jsonify(response.json())
    except requests.exceptions.RequestException as e:
        return jsonify({"error": str(e)}), 500


@app.route('/ask', methods=['POST'])
def ask_question():
    """提问接口"""
    data = request.get_json()
    if not data or 'message' not in data:
        return jsonify({"error": "Missing 'message' in request"}), 400

    # 构造请求
    payload = {
        "message": data['message'],
        "mode": data.get('mode', 'query'),
        "context": {
            "document_ids": data.get('document_ids', [])
        }
    }

    try:
        response = requests.post(
            url=f"{API_BASE}/prompt",
            headers={
                "Authorization": API_KEY,
                "Content-Type": "application/json"
            },
            json=payload,
            timeout=15
        )
        response.raise_for_status()
        return jsonify(response.json())
    except requests.exceptions.RequestException as e:
        return jsonify({"error": str(e)}), 500


# 简易测试页面
@app.route('/')
def test_page():
    return render_template_string(
        '''
        <h3>文件上传测试</h3>
        <form method="post" action="/upload" enctype="multipart/form-data">
            <input type="file" name="file">
            <input type="text" name="name" placeholder="文档名称">
            <button type="submit">上传</button>
        </form>

        <h3>提问测试</h3>
        <form method="post" action="/ask">
            <input type="text" name="message" placeholder="输入问题" style="width:300px;">
            <input type="text" name="document_ids" placeholder="文档ID (逗号分隔)">
            <button type="submit">提问</button>
        </form>
        ''')


if __name__ == '__main__':
    app.run(host='0.0.0.0', port=6789, debug=True)
